#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
股票历史数据管理器
功能：
1. 从公开数据源获取A股历史行情数据
2. 计算技术指标（均线、成交量均线等）
3. 数据存储到SQLite数据库（按年分表）
4. 支持数据更新、查询和清理
"""

import akshare as ak
import requests
import time
import pandas as pd
import logging
import sqlite3
import os
import random
import requests
import urllib3
from tqdm import tqdm
from datetime import datetime, timedelta
from typing import Dict, Optional, List
import numpy as np
from path_helper import PathManager
from trade_date.trade_date_manager import TradeDateManager

# 配置日志，设置日志级别和格式
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


# 禁用不安全的请求警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


class StockHistoryManager:
    """股票历史数据管理器

    特性：
    - 自动维护数据库表结构（按年分表）
    - 支持批量数据更新（带进度条显示）
    - 数据缓存机制减少网络请求
    - 完整的异常处理和日志记录
    - 增强的网络错误处理（重试、限流、代理）

    数据库设计：
    - stock_basic_info 表存储股票元数据
    - stock_data_YYYY 年分表存储行情数据
    """

    def __init__(self,
                 db_name: str = 'stock_history_data.db',
                 max_retries: int = 5,
                 retry_delay: int = 10,
                 request_delay: float = 0.5,
                 proxy_list: List[str] = None):
        """ 股票历史数据收集器初始化 """
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.request_delay = request_delay
        self.proxy_list = proxy_list or []
        self.last_request_time = None

        self.trade_calendar = TradeDateManager()  # 初始化交易日历管理器
        pm = PathManager()  # 初始化路径管理器
        self.db_name = pm.get_database_path(db_name)  # 获取数据库的完整路径
        self.create_database()  # 创建数据库
        self.configure_database()  # 配置数据库参数

        # 数据缓存字典和缓存过期时间设置
        self._query_cache = {}
        self.cache_expiry = timedelta(minutes=30)  # 设置缓存过期时间为30分钟

    def create_database(self) -> None:
        """ 创建 SQLite 数据库和相关数据表 """
        # 检查数据库是否存在
        if not os.path.exists(self.db_name):
            logging.info(f"数据库 {self.db_name} 不存在，正在创建。")
            try:
                with sqlite3.connect(self.db_name) as conn:
                    self._create_tables(conn)  # 创建所需表格
                    conn.commit()  # 提交创建事务
                    logging.info("数据库初始化完成")
            except sqlite3.Error as e:
                logging.critical(f"数据库创建失败: {e}")
                if os.path.exists(self.db_name):
                    os.remove(self.db_name)  # 如果创建失败，删除数据库文件
                raise
        else:
            logging.info("数据库已存在，跳过初始化")  # 数据库已存在，跳过创建

    def _create_tables(self, conn: sqlite3.Connection) -> None:
        """ 创建股票基本信息表和按年分片的行情数据表 """
        # 创建股票基本信息表
        conn.execute('''
            CREATE TABLE IF NOT EXISTS stock_basic_info (
                stock_code TEXT PRIMARY KEY,       -- 股票代码
                stock_name TEXT,                   -- 股票名称
                total_market_value REAL,           -- 总市值
                circulating_market_value REAL,      -- 流通市值
                industry TEXT,                     -- 所属行业
                listing_date TEXT,                 -- 上市日期
                total_shares INTEGER,              -- 总股本
                circulating_shares INTEGER,         -- 流通股本
                rating_agency TEXT,                -- 评级机构
                rating REAL,                       -- 评级
                rating_quarter TEXT,               -- 评级季度
                market_type TEXT,                  -- 市场类型
                last_updated TEXT                   -- 最后更新时间
            )
        ''')
        current_year = datetime.now().year  # 获取当前年份
        # 创建每年的股票数据表
        for year in range(2007, current_year + 1):
            self._create_year_table(conn, year)

    def _create_year_table(self, conn: sqlite3.Connection, year: int) -> None:
        """创建指定年份的数据表

        Args:
            conn: 数据库连接
            year: 年份
        """
        table_name = f'stock_data_{year}'  # 根据年份生成表名
        conn.execute(f'''
            CREATE TABLE IF NOT EXISTS {table_name} (
                date TEXT NOT NULL,                   -- 交易日期
                symbol TEXT NOT NULL,                 -- 股票代码
                open REAL,                            -- 开盘价
                close REAL,                           -- 收盘价
                high REAL,                            -- 最高价
                low REAL,                             -- 最低价
                volume REAL,                          -- 成交量
                turnover REAL,                        -- 成交额
                amplitude REAL,                       -- 振幅
                change_pct REAL,                      -- 涨跌幅
                change_amt REAL,                      -- 涨跌额
                turnover_rate REAL,                   -- 换手率
                ma5 REAL,                             -- 5日均线
                ma10 REAL,                            -- 10日均线
                ma20 REAL,                            -- 20日均线
                ma30 REAL,                            -- 30日均线
                ma60 REAL,                            -- 60日均线
                ma120 REAL,                           -- 120日均线
                vol_ma5 REAL,                         -- 5日成交量均线
                vol_ma10 REAL,                        -- 10日成交量均线
                vol_ma20 REAL,                        -- 20日成交量均线
                vol_ma30 REAL,                        -- 30日成交量均线
                vol_ma60 REAL,                        -- 60日成交量均线
                vol_ma120 REAL,                       -- 120日成交量均线
                PRIMARY KEY (date, symbol)           -- 复合主键
            ) WITHOUT ROWID
        ''')
        # 创建索引以加速查询
        conn.execute(f'CREATE INDEX IF NOT EXISTS idx_{year}_symbol ON {table_name}(symbol)')
        conn.execute(f'CREATE INDEX IF NOT EXISTS idx_{year}_date ON {table_name}(date)')
        conn.execute(f'CREATE INDEX IF NOT EXISTS idx_{year}_symbol_date ON {table_name}(symbol, date)')

    def configure_database(self) -> None:
        """配置数据库参数以优化性能"""
        try:
            with sqlite3.connect(self.db_name) as conn:
                # 使用WAL模式提升并发性能
                conn.execute("PRAGMA journal_mode = WAL;")
                conn.execute("PRAGMA synchronous = NORMAL;")  # 设置同步模式
                conn.execute("PRAGMA cache_size = -20000;")  # 20MB缓存
                conn.execute("PRAGMA mmap_size = 1073741824;")  # 1GB内存映射
                conn.execute("PRAGMA temp_store = MEMORY;")  # 使用内存存储临时表
                conn.execute("PRAGMA busy_timeout = 30000;")  # 设置超时为30秒
        except sqlite3.Error as e:
            logging.error(f"数据库配置失败: {e}")  # 记录配置失败的错误信息

    def _request_limiter(self):
        """控制请求频率，避免过快请求"""
        if self.last_request_time:
            elapsed = (datetime.now() - self.last_request_time).total_seconds()
            if elapsed < self.request_delay:
                sleep_time = self.request_delay - elapsed
                time.sleep(sleep_time)
        self.last_request_time = datetime.now()

    def _get_proxy(self) -> Optional[dict]:
        """随机获取一个代理（如果有配置）"""
        if self.proxy_list:
            proxy = random.choice(self.proxy_list)
            return {"http": proxy, "https": proxy}
        return None

    def _safe_akshare_call(self, func, *args, **kwargs):
        """安全调用akshare函数，带重试和错误处理

        Args:
            func: akshare函数
            *args: 位置参数
            **kwargs: 关键字参数

        Returns:
            函数执行结果或空DataFrame
        """
        # 可选代理设置（如果配置了代理列表）
        if self.proxy_list:
            proxies = self._get_proxy()
            if proxies:
                kwargs["proxies"] = proxies

        for attempt in range(self.max_retries):
            try:
                self._request_limiter()
                return func(*args, **kwargs)
            except requests.exceptions.HTTPError as e:
                logging.error(f"HTTP错误: {e}")  # 处理HTTP错误
                break
            except (requests.exceptions.ConnectionError,
                    requests.exceptions.Timeout,
                    requests.exceptions.TooManyRedirects) as e:
                logging.warning(f"网络错误 ({e.__class__.__name__})，将在 {self.retry_delay} 秒后重试...")
                time.sleep(self.retry_delay)
            except Exception as e:
                logging.error(f"调用 {func.__name__} 失败: {e}")
                if attempt < self.max_retries - 1:
                    time.sleep(self.retry_delay * (attempt + 1))  # 指数退避
                else:
                    break
        return pd.DataFrame()  # 返回空DataFrame

    def fetch_historical_data(self, symbol: str, start_date: str, end_date: str, adjust: str = "qfq") -> pd.DataFrame:
        """从数据源获取单只股票历史数据（带缓存机制）

        Args:
            symbol: 股票代码（如：000001）
            start_date: 开始日期（YYYYMMDD）
            end_date: 结束日期（YYYYMMDD）
            adjust: 复权类型（qfq:前复权，hfq:后复权，空:不复权）

        Returns:
            pd.DataFrame: 包含历史数据的DataFrame
        """
        # 参数校验
        try:
            datetime.strptime(start_date, "%Y%m%d")
            datetime.strptime(end_date, "%Y%m%d")
        except ValueError:
            logging.error("日期格式应为 YYYYMMDD")
            return pd.DataFrame()

        cache_key = f"{symbol}_{start_date}_{end_date}_{adjust}"  # 缓存键生成
        # 检查缓存
        if cache_key in self._query_cache:
            cached_data, timestamp = self._query_cache[cache_key]
            if datetime.now() - timestamp < self.cache_expiry:
                return cached_data.copy()  # 返回缓存数据的副本

        # 安全调用akshare
        data = self._safe_akshare_call(
            ak.stock_zh_a_hist,
            symbol=symbol,
            start_date=start_date,
            end_date=end_date,
            adjust=adjust
        )

        if data.empty:  # 检查数据是否为空
            logging.warning(f"未获取到 {symbol} 的历史数据")
            return pd.DataFrame()  # 返回空的 DataFrame

        # 更新缓存
        self._query_cache[cache_key] = (data.copy(), datetime.now())
        return data  # 返回获取到的数据

    def calculate_moving_averages(self, data: pd.DataFrame) -> pd.DataFrame:
        """计算移动平均指标

        计算包括：
        - 价格均线（5/10/20/30/60/120日）
        - 成交量均线（5/10/20/30/60/120日），单位为万手

        Args:
            data: 原始行情数据DataFrame

        Returns:
            添加了均线指标的新DataFrame
        """
        # 检查数据完整性
        if data.empty or '收盘' not in data.columns or '成交量' not in data.columns:
            logging.warning("无法计算移动平均线 - 数据为空或缺少必要列")
            return data  # 返回原始数据

        # 计算价格的移动平均线
        price_windows = [5, 10, 20, 30, 60, 120]
        for window in price_windows:
            data[f'ma{window}'] = (
                data['收盘']
                .rolling(window=window, min_periods=1)  # 计算移动平均
                .mean()
                .round(2)  # 保留两位小数
            )

        # 计算成交量的移动平均线，转换为万单位，保留一位小数
        volume_windows = [5, 10, 20, 30, 60, 120]
        for window in volume_windows:
            data[f'vol_ma{window}'] = (
                data['成交量']
                .rolling(window=window, min_periods=1)  # 计算移动平均
                .mean()
                .div(10000)  # 转换为万单位
                .round(1)  # 保留一位小数
            )

        return data  # 返回计算后的数据

    def upsert_historical_data(self, historical_data: pd.DataFrame) -> int:
        """ 更新或插入历史数据到数据库 """
        if historical_data.empty:
            logging.warning("尝试保存空数据")
            return 0  # 返回0条插入记录

        inserted_rows = 0  # 统计成功插入的行数
        try:
            with sqlite3.connect(self.db_name) as conn:
                conn.execute("BEGIN IMMEDIATE")  # 开始事务
                for _, row in historical_data.iterrows():
                    date = row['日期']  # 获取日期
                    symbol = row['股票代码']  # 获取股票代码
                    date_str = date.strftime('%Y-%m-%d')  # 格式化日期为字符串
                    exists_query = f'SELECT COUNT(1) FROM stock_data_{date.year} WHERE date = ? AND symbol = ?'
                    logging.info(f"查询: {exists_query} 参数: {date_str}, {symbol}")

                    # 检查数据是否存在
                    exists = conn.execute(exists_query, (date_str, symbol)).fetchone()[0]

                    if exists:
                        self._update_historical_data(conn, row, date, symbol)  # 更新数据
                        logging.info(f"更新数据: {symbol} 日期: {date_str}")
                    else:
                        self._insert_historical_data(conn, row, date, symbol)  # 插入数据
                        logging.info(f"插入数据: {symbol} 日期: {date_str}")

                    inserted_rows += 1  # 统计插入的行数

                conn.commit()  # 提交事务
                logging.info(f"成功插入或更新 {inserted_rows} 条记录")

        except sqlite3.Error as e:
            logging.error(f"数据库错误: {e}")  # 记录数据库错误
        except Exception as e:
            logging.error(f"发生错误: {e}")  # 记录其他可能的异常

        return inserted_rows  # 返回插入的行数

    def _update_historical_data(self, conn: sqlite3.Connection, row: pd.Series, date: str, symbol: str) -> None:
        """ 更新历史数据 """
        update_sql = f'''
            UPDATE stock_data_{date.year}
            SET open = ?, close = ?, high = ?, low = ?, volume = ?, 
                turnover = ?, amplitude = ?, change_pct = ?, change_amt = ?, 
                turnover_rate = ?, ma5 = ?, ma10 = ?, ma20 = ?, ma30 = ?, 
                ma60 = ?, ma120 = ?, vol_ma5 = ?, vol_ma10 = ?, vol_ma20 = ?, 
                vol_ma30 = ?, vol_ma60 = ?, vol_ma120 = ?
            WHERE date = ? AND symbol = ?
        '''

        conn.execute(update_sql, (
            row.get('开盘', 0.0), row.get('收盘', 0.0),  # 开盘和收盘价
            row.get('最高', 0.0), row.get('最低', 0.0),  # 最高和最低价
            row.get('成交量', 0), row.get('成交额', 0.0),  # 成交量和成交额
            row.get('振幅', 0.0), row.get('涨跌幅', 0.0),  # 振幅和涨跌幅
            row.get('涨跌额', 0.0), row.get('换手率', 0.0),  # 涨跌额和换手率
            row.get('ma5', 0.0), row.get('ma10', 0.0),  # 移动平均
            row.get('ma20', 0.0), row.get('ma30', 0.0),  # 移动平均
            row.get('ma60', 0.0), row.get('ma120', 0.0),  # 移动平均
            row.get('vol_ma5', 0.0), row.get('vol_ma10', 0.0),  # 成交量的移动平均
            row.get('vol_ma20', 0.0), row.get('vol_ma30', 0.0),  # 成交量的移动平均
            row.get('vol_ma60', 0.0), row.get('vol_ma120', 0.0),  # 成交量的移动平均
            date, symbol  # 更新条件
        ))

    def _insert_historical_data(self, conn: sqlite3.Connection, row: pd.Series, date: datetime, symbol: str) -> None:
        """ 插入历史数据 - 优化版本 """
        try:
            # 确保使用正确的日期格式
            date_str = date.strftime('%Y-%m-%d')

            # 获取字段值，带默认值处理
            def get_value(key, default=0.0):
                try:
                    return row[key]
                except KeyError:
                    logging.warning(f"字段 '{key}' 缺失，使用默认值 {default}")
                    return default

            # 准备数据
            row_data = (
                date_str, symbol,
                get_value('开盘'), get_value('收盘'),
                get_value('最高'), get_value('最低'),
                get_value('成交量', 0), get_value('成交额', 0.0),
                get_value('振幅', 0.0), get_value('涨跌幅', 0.0),
                get_value('涨跌额', 0.0), get_value('换手率', 0.0),
                get_value('ma5', 0.0), get_value('ma10', 0.0),
                get_value('ma20', 0.0), get_value('ma30', 0.0),
                get_value('ma60', 0.0), get_value('ma120', 0.0),
                get_value('vol_ma5', 0.0), get_value('vol_ma10', 0.0),
                get_value('vol_ma20', 0.0), get_value('vol_ma30', 0.0),
                get_value('vol_ma60', 0.0), get_value('vol_ma120', 0.0)
            )

            # SQL语句 - 确保字段数量和顺序匹配
            insert_sql = f'''
                INSERT INTO stock_data_{date.year} (
                    date, symbol, open, close, high, low, 
                    volume, turnover, amplitude, change_pct, 
                    change_amt, turnover_rate, ma5, ma10, 
                    ma20, ma30, ma60, ma120, vol_ma5, 
                    vol_ma10, vol_ma20, vol_ma30, vol_ma60, 
                    vol_ma120
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 
                          ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 
                          ?, ?, ?, ?)
            '''

            # 执行插入
            conn.execute(insert_sql, row_data)
            logging.info(f"成功插入数据: {symbol} {date_str}")

        except sqlite3.Error as e:
            logging.error(f"插入数据失败: {e}")
            logging.error(f"SQL: {insert_sql}")
            logging.error(f"数据: {row_data}")
            raise

    def fetch_combined_data(self, batch_size: int = 50) -> int:
        """主数据采集流程（批量处理）

        Args:
            batch_size: 每批处理的股票数量，默认50

        Returns:
            int: 成功处理股票数量
        """
        total_inserted = 0  # 统计成功插入的记录数

        try:
            logging.info("开始获取股票数据...")
            # 获取全量股票代码和ESG数据
            stock_info_a_code_name_df = self._safe_akshare_call(ak.stock_info_a_code_name)  # 股票代码和名称
            stock_esg_zd_sina_df = self._safe_akshare_call(ak.stock_esg_zd_sina)  # 获取ESG评分数据

            if stock_info_a_code_name_df.empty or stock_esg_zd_sina_df.empty:
                logging.error("获取股票基本信息失败，无法继续")
                return 0

            # 预处理ESG数据
            esg_data_dict = (
                stock_esg_zd_sina_df[["股票代码", "ESG评分", "评分日期"]]
                .set_index('股票代码')
                .to_dict(orient='index')
            )
            esg_data_dict = {k.split('.')[0]: v for k, v in esg_data_dict.items()}

            # 分批次处理
            stock_list = stock_info_a_code_name_df.to_dict('records')  # 将股票信息转换为字典列表
            total_batches = (len(stock_list) + batch_size - 1) // batch_size  # 计算总批次

            # 使用tqdm创建进度条
            for batch_idx in tqdm(range(total_batches), desc="🚀 处理股票进度", unit="batch"):
                start_idx = batch_idx * batch_size  # 当前批次起始索引
                end_idx = min((batch_idx + 1) * batch_size, len(stock_list))  # 当前批次结束索引
                batch_items = stock_list[start_idx:end_idx]  # 当前批次的股票信息

                # 批量处理逻辑
                with sqlite3.connect(self.db_name) as conn:
                    conn.execute("BEGIN IMMEDIATE")  # 开始事务
                    try:
                        for item in batch_items:
                            code = item['code']  # 股票代码
                            name = item['name']  # 股票名称

                            # 1. 更新基本信息
                            self._upsert_stock_basic_info(
                                conn=conn,
                                stock_code=code,
                                stock_name=name,
                                esg_info=esg_data_dict.get(code, {})  # 获取ESG信息
                            )

                            # 2. 获取历史行情数据
                            start_date = self._determine_start_date(code)  # 确定开始日期
                            raw_data = self.fetch_historical_data(
                                symbol=code,
                                start_date=start_date,
                                end_date=datetime.now().strftime("%Y%m%d")  # 获取当前日期作为结束日期
                            )

                            if raw_data.empty:
                                continue  # 如果没有数据，则跳过

                            # 3. 计算技术指标
                            processed_data = self.calculate_moving_averages(raw_data)

                            # 4. 补充元数据
                            processed_data['股票名称'] = name  # 添加股票名称
                            processed_data['市场类型'] = self.get_stock_exchange(code)  # 获取市场类型
                            esg_info = esg_data_dict.get(code, {})
                            processed_data['ESG评分'] = esg_info.get('ESG评分', np.nan)  # 添加ESG评分
                            processed_data['评分机构'] = '秩鼎' if esg_info else '无'  # 添加评分机构

                            # 5. 批量写入数据库
                            self._batch_upsert_data(conn, processed_data, code)

                        conn.commit()  # 提交事务
                        total_inserted += len(batch_items)  # 更新插入计数
                    except Exception as e:
                        conn.rollback()  # 回滚事务
                        logging.error(f"批次{batch_idx}处理失败: {str(e)}")  # 记录处理失败的错误信息

                # 批次处理间隔
                if batch_idx < total_batches - 1:
                    time.sleep(15)  # 每批处理间隔15秒

            logging.info(f"成功处理 {total_inserted}/{len(stock_list)} 条股票数据")
            return total_inserted  # 返回处理结果

        except Exception as e:
            logging.error(f"数据获取流程异常: {str(e)}")  # 记录异常信息
            return total_inserted  # 返回处理结果

    def _batch_upsert_data(self, conn: sqlite3.Connection, data: pd.DataFrame, symbol: str) -> None:
        """ 批量更新/插入行情数据（使用同一个连接） """
        if data.empty:
            return  # 如果数据为空，直接返回

        # 获取第一个日期用于确定年份分片
        first_date = pd.to_datetime(data.iloc[0]['日期'])  # 获取第一个日期
        table_name = f"stock_data_{first_date.year}"  # 根据年份确定表名

        # 准备批量数据
        records = []
        for _, row in data.iterrows():
            records.append((
                row['日期'].strftime('%Y-%m-%d'),  # 日期
                symbol,  # 股票代码
                row['开盘'], row['收盘'], row['最高'], row['最低'],  # 价格数据
                row['成交量'], row['成交额'], row['振幅'],  # 成交数据
                row['涨跌幅'], row['涨跌额'], row['换手率'],  # 涨跌数据
                row.get('ma5', 0), row.get('ma10', 0),  # 移动平均
                row.get('ma20', 0), row.get('ma30', 0),  # 移动平均
                row.get('ma60', 0), row.get('ma120', 0),  # 移动平均
                row.get('vol_ma5', 0), row.get('vol_ma10', 0),  # 成交量的移动平均
                row.get('vol_ma20', 0), row.get('vol_ma30', 0),  # 成交量的移动平均
                row.get('vol_ma60', 0), row.get('vol_ma120', 0)  # 成交量的移动平均
            ))

        # 使用批量插入ON CONFLICT UPDATE
        insert_sql = f"""
            INSERT INTO {table_name} VALUES (
                ?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?
            )
            ON CONFLICT(date, symbol) DO UPDATE SET
                open=excluded.open,
                close=excluded.close,
                high=excluded.high,
                low=excluded.low,
                volume=excluded.volume,
                turnover=excluded.turnover,
                amplitude=excluded.amplitude,
                change_pct=excluded.change_pct,
                change_amt=excluded.change_amt,
                turnover_rate=excluded.turnover_rate,
                ma5=excluded.ma5,
                ma10=excluded.ma10,
                ma20=excluded.ma20,
                ma30=excluded.ma30,
                ma60=excluded.ma60,
                ma120=excluded.ma120,
                vol_ma5=excluded.vol_ma5,
                vol_ma10=excluded.vol_ma10,
                vol_ma20=excluded.vol_ma20,
                vol_ma30=excluded.vol_ma30,
                vol_ma60=excluded.vol_ma60,
                vol_ma120=excluded.vol_ma120
            """
        conn.executemany(insert_sql, records)  # 批量执行插入

    def _upsert_stock_basic_info(self, conn: sqlite3.Connection, stock_code: str,
                                 stock_name: str, esg_info: dict) -> None:
        """ 使用现有连接更新股票基本信息 """
        try:
            # 获取个股详细信息
            info_df = self._safe_akshare_call(
                ak.stock_individual_info_em,
                symbol=stock_code
            )

            if info_df.empty:
                logging.warning(f"无法获取股票 {stock_code} 的详细信息")
                return

            info_dict = info_df.set_index('item')['value'].to_dict()  # 转换为字典格式

            # 构造插入数据
            insert_data = (
                stock_code,
                stock_name,
                info_dict.get('总市值', 0),
                info_dict.get('流通市值', 0),
                info_dict.get('行业', '未知'),
                info_dict.get('上市时间', '1970-01-01'),
                info_dict.get('总股本', 0),
                info_dict.get('流通股', 0),
                '秩鼎' if esg_info else '无',
                esg_info.get('ESG评分', 0),
                esg_info.get('评分日期', '未评级'),
                self.get_stock_exchange(stock_code),
                datetime.now().strftime("%Y-%m-%d %H:%M:%S")  # 最后更新时间
            )

            # 执行插入或更新操作
            conn.execute('''
                INSERT INTO stock_basic_info VALUES (
                    ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
                )
                ON CONFLICT(stock_code) DO UPDATE SET
                        stock_name = excluded.stock_name,
                        total_market_value = excluded.total_market_value,
                        circulating_market_value = excluded.circulating_market_value,
                        industry = excluded.industry,
                        listing_date = excluded.listing_date,
                        total_shares = excluded.total_shares,
                        circulating_shares = excluded.circulating_shares,
                        rating_agency = excluded.rating_agency,
                        rating = excluded.rating,
                        rating_quarter = excluded.rating_quarter,
                        market_type = excluded.market_type,
                        last_updated = excluded.last_updated
            ''', insert_data)

        except Exception as e:
            logging.error(f"股票{stock_code}基本信息更新失败: {str(e)}")  # 记录更新失败的错误信息

    def _determine_start_date(self, stock_code: str) -> str:
        """ 确定数据开始日期（带容错处理） """
        try:
            # 假设所有股票的开始日期为2007-01-01
            return "20070101"  # 返回默认开始日期
        except Exception as e:
            logging.warning(f"日期解析异常 [{stock_code}]：{str(e)}")  # 记录日期解析异常
            return "20070101"  # 返回默认开始日期

    def get_stock_exchange(self, code: str) -> str:
        """ 根据股票代码判断所属交易所 """
        code_str = str(code).zfill(6)  # 确保股票代码为6位

        if code_str.startswith('60') or code_str.startswith('688'):
            return 'sh'  # 上海证券交易所
        elif code_str[:3] in ['000', '001', '002', '003'] or code_str.startswith('300'):
            return 'sz'  # 深圳证券交易所
        elif code_str[:2] in ['43', '83', '87', '88']:
            return 'bj'  # 北京证券交易所
        else:
            return 'unknown'  # 未知市场类型

    def query_stock_data(self, symbol: str, start_date: Optional[str] = None,
                         end_date: Optional[str] = None) -> pd.DataFrame:
        """查询股票历史数据（带基本面信息）

        Args:
            symbol: 股票代码
            start_date: 开始日期（YYYY-MM-DD）
            end_date: 结束日期（YYYY-MM-DD）

        Returns:
            包含以下字段的DataFrame：
            - 行情数据（开高低收等）
            - 技术指标（各周期均线）
            - 基本面信息（市值、行业等）
            - ESG评分信息
        """
        start_date = start_date or "2007-01-01"  # 默认开始日期
        end_date = end_date or datetime.now().strftime("%Y-%m-%d")  # 默认结束日期

        start_year = pd.to_datetime(start_date).year  # 获取开始年份
        end_year = pd.to_datetime(end_date).year  # 获取结束年份

        results = []  # 存储查询结果的列表
        with sqlite3.connect(self.db_name) as conn:
            # 遍历每个年份进行查询
            for year in range(start_year, end_year + 1):
                table_name = f"stock_data_{year}"  # 生成表名
                try:
                    query = f'''
                    SELECT sd.date, sd.symbol, sd.open, sd.close, sd.high, sd.low, sd.volume,
                           sd.turnover, sd.amplitude, sd.change_pct, sd.change_amt, sd.turnover_rate,
                           sd.ma5, sd.ma10, sd.ma20, sd.ma30, sd.ma60, sd.ma120,
                           sd.vol_ma5, sd.vol_ma10, sd.vol_ma20, sd.vol_ma30, sd.vol_ma60, sd.vol_ma120,
                           bi.stock_name, bi.total_market_value, bi.circulating_market_value,
                           bi.industry, bi.listing_date, bi.total_shares, bi.circulating_shares,
                           bi.rating_agency, bi.rating, bi.rating_quarter, bi.market_type
                    FROM {table_name} sd
                    LEFT JOIN stock_basic_info bi ON sd.symbol = bi.stock_code
                    WHERE sd.symbol = ? AND sd.date BETWEEN ? AND ?
                    '''

                    year_data = conn.execute(query, (symbol, start_date, end_date)).fetchall()  # 执行查询
                    results.extend(year_data)  # 将查询结果添加到结果列表
                except sqlite3.OperationalError:
                    logging.warning(f"表 {table_name} 不存在，跳过查询。")  # 如果表不存在，记录警告
                    continue

        # 将查询结果转换为 DataFrame
        df = pd.DataFrame(results, columns=[
            'date', 'symbol', 'open', 'close', 'high', 'low', 'volume',
            'turnover', 'amplitude', 'change_pct', 'change_amt', 'turnover_rate',
            'ma5', 'ma10', 'ma20', 'ma30', 'ma60', 'ma120',
            'vol_ma5', 'vol_ma10', 'vol_ma20', 'vol_ma30', 'vol_ma60', 'vol_ma120',
            'stock_name', 'total_market_value', 'circulating_market_value',
            'industry', 'listing_date', 'total_shares', 'circulating_shares',
            'rating_agency', 'rating', 'rating_quarter', 'market_type'
        ])

        return df  # 返回查询结果

    def get_stock_codes_for_year(self, year: int, market_types: list = None,
                                 excluded_prefixes: list = None) -> pd.DataFrame:
        """ 获取指定年份的所有股票代码，并根据市场类型和排除的前缀进行筛选 """
        table_name = f"stock_basic_info"  # 股票基本信息表名

        # 设置默认值
        if market_types is None:
            market_types = ['sz', 'sh']  # 默认市场类型
        if excluded_prefixes is None:
            excluded_prefixes = ['688', '300']  # 默认排除前缀

        # 构建市场类型的 SQL 查询条件
        market_condition = "market_type IN ({})".format(", ".join(f"'{mt}'" for mt in market_types))

        # 构建排除前缀的 SQL 查询条件
        exclusion_conditions = " AND ".join(f"stock_code NOT LIKE '{prefix}%'" for prefix in excluded_prefixes)

        # 完整的 SQL 查询语句
        query = f'''
            SELECT DISTINCT stock_code 
            FROM {table_name}
            WHERE {market_condition} 
            AND {exclusion_conditions};
        '''

        with sqlite3.connect(self.db_name) as conn:
            try:
                stock_codes = conn.execute(query).fetchall()  # 执行查询并获取股票代码
                return pd.DataFrame(stock_codes, columns=['stock_code'])  # 返回股票代码的 DataFrame
            except sqlite3.OperationalError:
                logging.warning(f"表 {table_name} 不存在，无法查询股票代码。")  # 如果表不存在，记录警告
                return pd.DataFrame()  # 返回空的 DataFrame

    def delete_database(self) -> None:
        """ 删除数据库文件 """
        try:
            if os.path.exists(self.db_name):
                os.remove(self.db_name)  # 删除数据库文件
                logging.info(f"数据库 {self.db_name} 已成功删除")
            else:
                logging.warning(f"数据库 {self.db_name} 不存在，无法删除")  # 如果数据库不存在，记录警告
        except Exception as e:
            logging.error(f"删除数据库失败: {e}")  # 记录删除失败的错误信息

    def clear_data_for_year(self, year: int) -> None:
        """ 检查并清空指定年份的所有数据（如果存在） """
        logging.info(f"检查 {year} 年的数据...")
        table_name = f"stock_data_{year}"  # 获取指定年份的表名
        with sqlite3.connect(self.db_name) as conn:
            try:
                # 查询指定年份数据是否存在
                query = f'SELECT COUNT(1) FROM {table_name}'  # 计数查询
                exists = conn.execute(query).fetchone()[0] > 0  # 检查数据是否存在

                if exists:
                    # 删除指定年份的所有数据
                    delete_query = f'DELETE FROM {table_name}'  # 删除查询
                    conn.execute(delete_query)  # 执行删除
                    logging.info(f"{year} 年的所有数据已清空。")  # 记录清空成功

                else:
                    logging.info(f"{year} 年的数据不存在，无需清空。")  # 记录数据不存在的情况
            except sqlite3.OperationalError:
                logging.warning(f"表 {table_name} 不存在，无法清空数据。")  # 如果表不存在，记录警告


if __name__ == "__main__":
    # 配置代理列表（如果需要）
    # proxies = [
    #     "http://proxy1.example.com:8080",
    #     "http://proxy2.example.com:8080",
    #     "http://proxy3.example.com:8080"
    # ]

    # 实例化 StockHistoryManager
    collector = StockHistoryManager(
        max_retries=5,  # 最大重试次数
        retry_delay=15,  # 重试延迟（秒）
        request_delay=1.0,  # 请求间隔（秒）
        # 代理列表
    )

    # 1. 检查并清空2025年的数据（如果存在）
    current_year = datetime.now().year  # 获取当前年份
    if current_year >= 2025:
        collector.clear_data_for_year(2025)  # 清空2025年数据

    # 2. 收集所有股票数据
    total_inserted = collector.fetch_combined_data(batch_size=30)  # 使用较小的批次大小
    logging.info(f"总共插入 {total_inserted} 条记录")

    # 3. 查询特定股票数据示例
    stock_code = "000001"  # 示例股票代码
    start_date = "2024-01-01"  # 示例开始日期
    end_date = "2025-12-31"  # 示例结束日期

    # 查询数据
    stock_data = collector.query_stock_data(stock_code, start_date, end_date)

    # 输出查询到的股票数据
    if not stock_data.empty:  # 检查数据是否为空
        logging.info(f"查询到股票 {stock_code} 的数据:")
        logging.info(stock_data.head())  # 只输出前几行
    else:
        logging.warning(f"未找到股票 {stock_code} 的数据")  # 记录未找到数据的情况